import os
import torch
import numpy as np
import matplotlib.pyplot as plt
import SimpleITK as sitk
from tqdm import tqdm
from datetime import datetime
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix
from sklearn.metrics import ConfusionMatrixDisplay
import pandas as pd
import torch.nn.functional as F
from typing import Dict, Tuple, Any, List

# 导入之前定义的配置和模型
from train import ModelConfig, UNetPlusPlus as UNet, LiverDataset, dice_coeff  # 使用正确的模型类名

# 扩展配置类
class EvalConfig(ModelConfig):  # 继承自ModelConfig
    # 测试数据路径
    TEST_IMAGE_DIR = "testImage"
    TEST_MASK_DIR = "testMask"
    
    # 输出路径
    EVAL_OUTPUT_DIR = "evaluation_output"
    RESULTS_CSV = os.path.join(EVAL_OUTPUT_DIR, "test_results.csv")
    CONFUSION_MATRIX_PATH = os.path.join(EVAL_OUTPUT_DIR, "confusion_matrix.png")
    SAMPLE_VIS_PATH = os.path.join(EVAL_OUTPUT_DIR, "sample_predictions.png")
    
    def __init__(self):
        super().__init__()
        os.makedirs(self.EVAL_OUTPUT_DIR, exist_ok=True)

def load_model(config: EvalConfig) -> UNet:
    """加载训练好的模型
    
    Args:
        config: 配置对象
        
    Returns:
        UNet: 加载的模型
    """
    try:
        model = UNet(
            in_channels=config.IN_CHANNELS,
            out_channels=config.OUT_CHANNELS
        ).to(config.device)
        
        # 加载训练好的权重
        model.load_state_dict(torch.load(config.MODEL_SAVE_PATH))
        model.eval()
        
        print(f"模型已从 {config.MODEL_SAVE_PATH} 加载")
        return model
    except Exception as e:
        raise RuntimeError(f"模型加载失败: {str(e)}")

def create_test_dataset(config: EvalConfig) -> DataLoader:
    """创建测试数据集加载器
    
    Args:
        config: 配置对象
        
    Returns:
        DataLoader: 测试集数据加载器
    """
    try:
        test_dataset = LiverDataset(
            config.TEST_IMAGE_DIR, 
            config.TEST_MASK_DIR,
            transform=None
        )
        
        test_loader = DataLoader(
            test_dataset,
            batch_size=config.BATCH_SIZE,
            shuffle=False,
            num_workers=config.NUM_WORKERS,
            pin_memory=True
        )
        
        print(f"测试集大小: {len(test_dataset)} 个样本")
        return test_loader
    except Exception as e:
        raise RuntimeError(f"创建测试数据集失败: {str(e)}")

# === 混淆矩阵更新函数 ===
def update_confusion_matrix(accumulator: np.ndarray, true_labels: np.ndarray, pred_labels: np.ndarray) -> np.ndarray:
    """更新混淆矩阵累加器
    
    Args:
        accumulator: 当前累加的混淆矩阵
        true_labels: 真实标签 (H, W) 或 (N, H, W)
        pred_labels: 预测标签 (H, W) 或 (N, H, W)
        
    Returns:
        np.ndarray: 更新后的混淆矩阵
    """
    # 展平数组
    true_flat = true_labels.flatten()
    pred_flat = pred_labels.flatten()
    
    # 计算当前batch的混淆矩阵
    batch_cm = confusion_matrix(true_flat, pred_flat, labels=[0, 1, 2])
    
    # 确保矩阵形状正确 (3x3)
    if batch_cm.shape != (3, 3):
        full_cm = np.zeros((3, 3), dtype=np.int64)
        min_row = min(batch_cm.shape[0], 3)
        min_col = min(batch_cm.shape[1], 3)
        full_cm[:min_row, :min_col] = batch_cm[:min_row, :min_col]
        batch_cm = full_cm
    
    return accumulator + batch_cm

# === IoU指标计算函数 ===
def calculate_iou_metrics(preds: np.ndarray, masks: np.ndarray) -> Tuple[int, int, int, int]:
    """计算肝脏和肿瘤的IoU指标
    
    Args:
        preds: 预测标签 (N, H, W) 或 (H, W)
        masks: 真实标签 (N, H, W) 或 (H, W)
        
    Returns:
        Tuple: (肝脏交集, 肝脏并集, 肿瘤交集, 肿瘤并集)
    """
    # 确保维度一致
    if preds.ndim == 3:
        preds = preds.reshape(-1)
        masks = masks.reshape(-1)
    
    # 计算肝脏指标
    liver_pred = (preds == 1)
    liver_true = (masks == 1)
    liver_intersection = np.sum(liver_pred & liver_true)
    liver_union = np.sum(liver_pred | liver_true)
    
    # 计算肿瘤指标
    tumor_pred = (preds == 2)
    tumor_true = (masks == 2)
    tumor_intersection = np.sum(tumor_pred & tumor_true)
    tumor_union = np.sum(tumor_pred | tumor_true)
    
    return liver_intersection, liver_union, tumor_intersection, tumor_union

def calculate_metrics(
    outputs: torch.Tensor, 
    masks: torch.Tensor
) -> Tuple[float, float, float, float, np.ndarray, np.ndarray]:
    """计算评估指标"""
    # 使用明确的返回值结构
    bg_dice, liver_dice, tumor_dice, mean_dice, small_tumor_dice = dice_coeff(outputs, masks)
    
    # 转换为类别预测
    probs = F.softmax(outputs, dim=1)
    preds = probs.argmax(dim=1).cpu().numpy()
    masks_np = masks.cpu().numpy()
    
    # 返回6个值，与调用点一致
    return liver_dice, tumor_dice, mean_dice, small_tumor_dice, preds, masks_np

def evaluate_model(
    model: UNet, 
    test_loader: DataLoader, 
    config: EvalConfig
) -> Dict[str, Any]:
    """评估模型在测试集上的性能
    
    Args:
        model: 待评估的模型
        test_loader: 测试集数据加载器
        config: 配置对象
        
    Returns:
        Dict: 包含评估结果的字典
    """
    device = config.device
    
    # 初始化存储变量
    metrics = {
        'liver_dice_sum': 0.0,
        'tumor_dice_sum': 0.0,
        'mean_dice_sum': 0.0,
        'small_tumor_dice_sum': 0.0,  # 小肿瘤Dice统计
        'batch_count': 0,
        'cm_accumulator': np.zeros((3, 3), dtype=np.int64),
        'liver_intersection': 0,
        'liver_union': 0,
        'tumor_intersection': 0,
        'tumor_union': 0,
        'sample_images': None,
        'sample_masks': None,
        'sample_preds': None
    }
    
    with torch.no_grad():
        for i, (images, masks) in enumerate(tqdm(test_loader, desc="评估测试集")):
            images = images.to(device)
            masks = masks.to(device)
            
            # 模型预测
            outputs = model(images)
            
            # 计算指标
            liver_dice, tumor_dice, mean_dice, small_tumor_dice, preds, masks_np = calculate_metrics(outputs, masks)
            
            # 更新指标
            metrics['liver_dice_sum'] += liver_dice
            metrics['tumor_dice_sum'] += tumor_dice
            metrics['mean_dice_sum'] += mean_dice
            metrics['small_tumor_dice_sum'] += small_tumor_dice  # 新增
            metrics['batch_count'] += 1
            
            # 保存第一批次样本用于可视化
            if i == 0:
                metrics['sample_images'] = images.cpu().numpy()
                metrics['sample_masks'] = masks_np
                metrics['sample_preds'] = preds
            
            # 更新混淆矩阵
            metrics['cm_accumulator'] = update_confusion_matrix(
                metrics['cm_accumulator'], masks_np, preds
            )
            
            # 更新IoU指标
            li, lu, ti, tu = calculate_iou_metrics(preds, masks_np)
            metrics['liver_intersection'] += li
            metrics['liver_union'] += lu
            metrics['tumor_intersection'] += ti
            metrics['tumor_union'] += tu
            
            # 及时释放内存
            del outputs, preds, masks_np
            torch.cuda.empty_cache()
    
    # 计算最终指标
    avg_liver_dice = metrics['liver_dice_sum'] / metrics['batch_count']
    avg_tumor_dice = metrics['tumor_dice_sum'] / metrics['batch_count']
    avg_mean_dice = metrics['mean_dice_sum'] / metrics['batch_count']
    avg_small_tumor_dice = metrics['small_tumor_dice_sum'] / metrics['batch_count']  # 新增
    avg_liver_iou = metrics['liver_intersection'] / (metrics['liver_union'] + 1e-8)
    avg_tumor_iou = metrics['tumor_intersection'] / (metrics['tumor_union'] + 1e-8)
    
    # 生成分类报告
    class_report = generate_classification_report(metrics['cm_accumulator'])
    
    # 打印汇总指标
    print("\n评估结果汇总:")
    print(f"肝脏平均Dice系数: {avg_liver_dice:.4f}")
    print(f"肿瘤平均Dice系数: {avg_tumor_dice:.4f}")
    print(f"总体平均Dice系数: {avg_mean_dice:.4f}")
    print(f"肝脏平均IoU: {avg_liver_iou:.4f}")
    print(f"肿瘤平均IoU: {avg_tumor_iou:.4f}")
    
    return {
        'avg_liver_dice': avg_liver_dice,
        'avg_tumor_dice': avg_tumor_dice,
        'avg_mean_dice': avg_mean_dice,
        'avg_small_tumor_dice': avg_small_tumor_dice,  # 新增
        'avg_liver_iou': avg_liver_iou,
        'avg_tumor_iou': avg_tumor_iou,
        'confusion_matrix': metrics['cm_accumulator'],
        'classification_report': class_report,
        'sample_images': metrics['sample_images'],
        'sample_masks': metrics['sample_masks'],
        'sample_preds': metrics['sample_preds'],
    }

def generate_classification_report(cm: np.ndarray) -> Dict[str, Dict]:
    """从混淆矩阵生成分类报告
    
    Args:
        cm: 混淆矩阵
        
    Returns:
        Dict: 分类报告字典
    """
    class_report = {}
    for idx, name in enumerate(['Background', 'Liver', 'Tumor']):
        if idx == 0:  # 跳过背景类
            continue
            
        tp = cm[idx, idx]
        fp = cm[:, idx].sum() - tp
        fn = cm[idx, :].sum() - tp
        tn = cm.sum() - (tp + fp + fn)
        
        precision = tp / (tp + fp + 1e-8)
        recall = tp / (tp + fn + 1e-8)
        f1 = 2 * (precision * recall) / (precision + recall + 1e-8)
        
        class_report[name] = {
            'precision': precision,
            'recall': recall,
            'f1-score': f1,
            'support': int(tp + fn)
        }
    
    return class_report

def visualize_confusion_matrix(cm: np.ndarray, config: EvalConfig) -> None:
    """可视化混淆矩阵并保存
    
    Args:
        cm: 混淆矩阵
        config: 配置对象
    """
    plt.figure(figsize=(10, 8))
    disp = ConfusionMatrixDisplay(
        confusion_matrix=cm,
        display_labels=['Background', 'Liver', 'Tumor']
    )
    disp.plot(cmap='Blues', values_format='d')
    plt.title('Confusion Matrix')
    plt.savefig(config.CONFUSION_MATRIX_PATH)
    plt.close()
    print(f"混淆矩阵已保存至: {config.CONFUSION_MATRIX_PATH}")

def visualize_sample_predictions(
    images: np.ndarray, 
    masks: np.ndarray, 
    preds: np.ndarray, 
    config: EvalConfig, 
    num_samples: int = 3
) -> None:
    """可视化样本的预测结果并保存
    
    Args:
        images: 图像数据
        masks: 真实标签
        preds: 预测结果
        config: 配置对象
        num_samples: 显示的样本数量
    """
    plt.figure(figsize=(15, 5 * num_samples))
    
    for i in range(min(num_samples, images.shape[0])):
        # 获取中间切片
        img_slice = images[i, 0]  # 取第一个通道
        mask_slice = masks[i]
        pred_slice = preds[i]
        
        # 原始图像
        plt.subplot(num_samples, 3, i * 3 + 1)
        plt.imshow(img_slice, cmap='gray')
        plt.title(f"样本 {i+1}: 原始CT")
        plt.axis('off')
        
        # 真实标签
        plt.subplot(num_samples, 3, i * 3 + 2)
        plt.imshow(mask_slice, cmap='jet', vmin=0, vmax=2)
        plt.title("真实标签")
        plt.axis('off')
        
        # 预测结果
        plt.subplot(num_samples, 3, i * 3 + 3)
        plt.imshow(pred_slice, cmap='jet', vmin=0, vmax=2)
        plt.title("预测结果")
        plt.axis('off')
    
    plt.tight_layout()
    plt.savefig(config.SAMPLE_VIS_PATH)
    plt.close()
    print(f"样本预测可视化已保存至: {config.SAMPLE_VIS_PATH}")

def generate_evaluation_report(results: dict, config: EvalConfig) -> str:
    """生成评估报告并保存
    
    Args:
        results: 评估结果字典
        config: 配置对象
        
    Returns:
        str: 评估报告内容
    """
    report = f"""
    {'='*50}
    肝脏与肝肿瘤分割模型评估报告
    {'='*50}
    
    评估时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
    测试集大小: {len(os.listdir(config.TEST_IMAGE_DIR))} 个切片
    
    关键指标:
    - 肝脏平均Dice系数: {results['avg_liver_dice']:.4f}
    - 肿瘤平均Dice系数: {results['avg_tumor_dice']:.4f}
    - 总体平均Dice系数: {results['avg_mean_dice']:.4f}
    - 肝脏平均IoU: {results['avg_liver_iou']:.4f}
    - 肿瘤平均IoU: {results['avg_tumor_iou']:.4f}
    
    详细结果请查看:
    - 样本预测可视化: {config.SAMPLE_VIS_PATH}
    - 混淆矩阵: {config.CONFUSION_MATRIX_PATH}
    - 完整测试结果: {config.RESULTS_CSV}
    """
    
    report_path = os.path.join(config.EVAL_OUTPUT_DIR, "evaluation_report.txt")
    with open(report_path, 'w') as f:
        f.write(report)
    
    print(f"评估报告已保存至: {report_path}")
    return report

def main() -> None:
    """主函数，执行评估流程"""
    start_time = datetime.now()
    print(f"评估开始时间: {start_time.strftime('%Y-%m-%d %H:%M:%S')}")
    
    try:
        # 初始化配置
        config = EvalConfig()
        
        # 加载模型
        model = load_model(config)
        
        # 创建测试集
        test_loader = create_test_dataset(config)
        
        # 评估模型
        results = evaluate_model(model, test_loader, config)
        
        # 可视化结果
        visualize_confusion_matrix(results['confusion_matrix'], config)
        visualize_sample_predictions(
            results['sample_images'], 
            results['sample_masks'], 
            results['sample_preds'], 
            config
        )
        
        # 生成评估报告
        generate_evaluation_report(results, config)
        
    except Exception as e:
        print(f"评估过程中出错: {str(e)}")
    
    end_time = datetime.now()
    duration = end_time - start_time
    print(f"评估完成! 总耗时: {duration}")

if __name__ == "__main__":
    main()